package edu.arizona.sista.learn.model.groups;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import edu.arizona.sista.learn.model.ParamInterval;
import edu.arizona.sista.learn.model.groups.BodyPart.BodyPartName;
import edu.arizona.sista.learn.util.Utils;
import edu.arizona.sista.learn.util.Utils.ActorName;
import edu.arizona.sista.learn.util.Utils.MovieName;

/**
 * A scene contains all information pertaining to an arbitrary
 * duration where activities can occcur (e.g. a video clip).
 * 
 * @author Anh Tran
 *
 */
public class Scene {
	
	private MovieName _name;
	private List<Group> _groups;
	private Map<String, Thing> _things;
	private Map<Set<Person>, List<Double>> _pairs;
	
	public Scene(MovieName name) {
		_name = name;
		_groups = new ArrayList<Group>();
		
		_things = new HashMap<String, Thing>();
		_pairs = new HashMap<Set<Person>, List<Double>>();
	}
	
	
	public MovieName getName() {
		return _name;
	}
	
	
	public List<Group> getGroups() {
		return _groups;
	}
	
	public Map<String, Thing> getThings() {
		return _things;
	}
	
	public Map<Set<Person>, List<Double>> getPairs() {
		return _pairs;
	}
	
	public List<Person> getActors() {
		List<Person> actors = new ArrayList<Person>();
		for (Thing a : _things.values()) {
			if (a instanceof Person) {
				actors.add((Person)a);
			}
		}
		return actors;
	}
	
	public Person getActor(ActorName name) {
		if (_things.containsKey(name.toString())) {
			return (Person)_things.get(name);
		}
		return null;
	}
	
	public Thing getThing(String id) {
		if (_things.containsKey(id)) {
			return _things.get(id);
		}
		return null;
	}
	
	public Thing getOrAddThing(String id) {
		Thing t;
		if (_things.containsKey(id)) {
			t = _things.get(id);
		} else { 
			t = new Thing(id);
			_things.put(id, t);
		}
		return t;
	}
	
	public Person getOrAddPerson(ActorName name) {
		Person p;
		String id = name.toString();
		if (_things.containsKey(id)) {
			p = (Person)_things.get(id);
		} else { 
			p = new Person(name.toString());
			_things.put(id, p);
		}
		return p;
	}
	
	public Thing getOrAddObject(String id) {
		Thing o;
		id = "Object" + id;
		if (_things.containsKey(id)) {
			o = _things.get(id);
		} else { 
			o = new Thing(id);
			_things.put(id, o);
		}
		return o;
	}
	
	public Thing getOrAddBodyPart(ActorName owner, BodyPartName part) {
		BodyPart b;
		String id = owner.toString() + "." + part.toString();
		if (_things.containsKey(id)) {
			b = (BodyPart)_things.get(id);
		} else { 
			b = new BodyPart(part, this.getOrAddPerson(owner));
			_things.put(id, b);
		}
		return b;
	}
	
	/**
	 * Load initial grouping configurations from file.
	 * @param dataFile Cluster data file.
	 */
	public void loadGroupConfigurations(File dataFile) {
		loadGroupConfigurations(dataFile, true);
	}
	
	public void loadGroupConfigurations(File dataFile, boolean out) {
		Map<String, Group> groupsMap = new HashMap<String, Group>();
		
		try {
			// Create BufferedReader to read csv file
			BufferedReader br = new BufferedReader(new FileReader(dataFile));
			String strLine = "";
			int prevFrame = -1, frame = 0;
			final String INVALID_GROUP_ID = "-1";
			
			// Skip first line
			strLine = br.readLine();

			// Read comma separated file line by line
			while ((strLine = br.readLine()) != null) {

				// Break comma separated line using ";"
				Set<String> gDone = new HashSet<String>();
				String[] tokens = strLine.split(";");
				if (tokens.length > 1) {
					
					// First column is frame number
					frame = Integer.parseInt(tokens[0].trim());
					assert prevFrame == -1 || frame == prevFrame+1 :
						"Missing cluster frame(s) in "
							+ dataFile.getName() + " between "
							+ prevFrame + " " + frame;
					
					// Subsequent columns are groups' (ID : [members]) pairs
					for (int i = 1; i < tokens.length; i++) {
						if (tokens[i].trim().length() == 0) continue;
						String[] groupStr = tokens[i].trim().split(":");
						assert groupStr.length == 2 : "Invalid group description: " + groupStr;
						String gID = groupStr[0].trim();
						if (!gID.equals(INVALID_GROUP_ID) && !gDone.contains(gID)) {
							String partsStr = groupStr[1].replaceAll("^\\[|\\]$", "");
							String[] parts = partsStr.split(",");
							Set<Person> members = new HashSet<Person>();
							for (String part : parts) {
								if (part.trim().length() == 0) continue;
								ActorName partName = Utils.getActorNameFromClusterID(_name, part.trim());
								members.add(this.getOrAddPerson(partName));
							}
							
							if (members.size() > 0) {
								// If we have a valid group ID and valid number of
								// participants, then formulate & update group.
								Group g;
								if (!groupsMap.containsKey(gID)) {
									g = new Group(gID, frame, frame+1, members);
									groupsMap.put(gID, g);
								} else {
									groupsMap.get(gID).extendLastGroupingVariation(1);
								}
								
								gDone.add(gID);
							}
						}
					}
				}
				
				prevFrame = frame;
			}

		} catch (Exception e) {
			System.out.println("Exception while reading csv file: " + e);
		}

		// Save parsed information to local variables
		_groups = new ArrayList<Group>(groupsMap.values());
		
		// Dump to file
		if (out)
			Utils.toFile(this.printGroups(), Utils.OUTDIR + _name + "-clustered.txt", false);
	}
	
	
	/**
	 * Load look-at data from file.
	 * @param dataFile Look-at data file.
	 */
	public void loadGazeData(File dataDir) {
		loadGazeData(dataDir, true);
	}
	
	public void loadGazeData(File dataDir, boolean out) {
		
		try {
			if (!dataDir.isDirectory())
				throw new RuntimeException("Invalid directory.");
			
			for (File dataFile : dataDir.listFiles()) {
				if (dataFile.getName().startsWith(".")) {
					continue;  // Ignore hidden files
				}
				
				ActorName sourceName = Utils.getActorNameFromTrackID(_name, dataFile.getName().trim());
				Person source = this.getOrAddPerson(sourceName);

				// Create BufferedReader to read csv file
				BufferedReader br = new BufferedReader(new FileReader(dataFile));
				String strLine = "";
				String[] targetTokens;
				BodyPartName targetPart;
				ActorName targetName;
				int prevFrame = -1, frame = 0;
				final String personType = "Person", objectType = "Object";
				List<ParamInterval<Thing>> currentIntervals = new ArrayList<ParamInterval<Thing>>();
				Set<Thing> targets = new HashSet<Thing>();

				// Read file line by line
				while ((strLine = br.readLine()) != null) {
					
					// Split line based on space
					String[] tokens = strLine.trim().split(" ");
					if (tokens.length >= 3) {
						
						// First column is frame number
						frame = Integer.parseInt(tokens[0].trim());

						// NEW 05/21: From now on, we'll enforce continuity
						// in frames. That is, we'll no longer interpolate
						// between missing frames.
						if (prevFrame != -1 && frame != prevFrame+1) {
							// If there are missing frames, then end all
							// previously rolling intervals.
							for (int i = currentIntervals.size()-1; i >= 0; i--) {
								source.addLookingInterval(currentIntervals.get(i));
								currentIntervals.remove(i);
							}
						}
						
						// Subsequent columns are look-at targets, read in quadruples.
						for (int i = 1; i < tokens.length; i += 4) {
							
							// First token in the quad is the target, we'll need to parse it.
							Thing target;
							targetTokens = tokens[i].trim().split("_");
							if (targetTokens[0].trim().equalsIgnoreCase(personType)) {
								targetName = Utils.getActorNameFromTrackID(_name, targetTokens[1].trim());
								targetPart = (targetTokens.length <= 2) ? null :
									BodyPartName.fromString(targetTokens[2].trim());
								if (targetPart != null) {
									target = this.getOrAddBodyPart(targetName, targetPart);
								} else {
									// looking at the full body (the person himself)
									target = this.getOrAddPerson(targetName);
								}
							} else if (targetTokens[0].trim().equalsIgnoreCase(objectType)) {
								target = this.getOrAddObject(targetTokens[1].trim());
							} else {
								// Shouldn't be here
								throw new RuntimeException("Unrecognizable target");
							}
							
							targets.add(target);
							
							// The second token in the pair is some score for that target.
							// We'll ignore that for now.
							// double score = new Double(t okens[i+1].trim());
						}
					}

					// Now that we parsed the targets, it's time to add to the intervals.
					// Note: We will need to interpolate between any missing frames for now
					// 		since we're working with spotty data.
					// For each rolling interval, we'll either extend it if the target is
					// still valid, or end it otherwise.
					for (int i = currentIntervals.size()-1; i >= 0; i--) {
						ParamInterval<Thing> interval = currentIntervals.get(i);
						if (targets.contains(interval.obj())) {
							// Extend the interval if the target is still valid.
							interval.interval().setEnd(frame+1);
							targets.remove(interval.obj());
						} else {
							// Otherwise, end it.
							source.addLookingInterval(interval);
							currentIntervals.remove(i);
						}
					}
					
					// For each new target, start a new rolling interval
					for (Thing target : targets) {
						assert (target != null) : "Invalid target!";
						currentIntervals.add(new ParamInterval<Thing>(frame, frame+1, target));
					}
					targets.clear();
					
					prevFrame = frame;
				}
				
				// By the time we finish reading the file, if there are any rolling intervals
				// left, then end them.
				for (int i = currentIntervals.size()-1; i >= 0; i--) {
					ParamInterval<Thing> interval = currentIntervals.get(i);
					interval.interval().setEnd(frame+1);
					source.addLookingInterval(interval);
					currentIntervals.remove(i);
				}
			}
			
		} catch (Exception e) {
			System.out.println("Exception while reading file: " + e);
		}

		// Dump to file
		if (out) {
			StringBuilder sb = new StringBuilder();
			for (Person p : this.getActors())
				sb.append(p.toStringLookAt() + "\n");
			System.out.println(sb.toString());
			Utils.toFile(sb, Utils.OUTDIR + _name + "-lookAt.txt", false);
		}
	}
	

	/**
	 * Load gestures data from file.
	 * @param dataFile Gestures data file.
	 */
	public void loadGesturesData(File dataDir) {
		loadGesturesData(dataDir, true);
	}
	
	public void loadGesturesData(File dataDir, boolean out) {
		
		try {
			if (!dataDir.isDirectory())
				throw new RuntimeException("Invalid directory.");
			
			for (File dataFile : dataDir.listFiles()) {
				if (dataFile.getName().startsWith(".") ||
						!dataFile.getName().endsWith(".txt")) {
					continue;  // Ignore hidden/irrelevant files
				}
				
				ActorName actorName = Utils.getActorNameFromTrackID(_name,
						dataFile.getName().trim().replace(".txt", ""));
				Person actor = this.getOrAddPerson(actorName);

				// Create BufferedReader to read csv file
				BufferedReader br = new BufferedReader(new FileReader(dataFile));
				String strLine = "";
				int prevFrame = -1, frame = 0;
				List<ParamInterval<String>> currentIntervals = new ArrayList<ParamInterval<String>>();
				Set<String> gestures = new HashSet<String>();

				// Skip first line
				strLine = br.readLine();
				
				// Read file line by line
				while ((strLine = br.readLine()) != null) {
					
					// Split line based on space
					String[] tokens = strLine.trim().split(";");
					if (tokens.length >= 2) {
						
						// First column is frame number
						frame = Integer.parseInt(tokens[0].trim());
						if (prevFrame != -1 && frame != prevFrame+1)
							System.err.println("Missing gestures frame(s) in "
									+ dataFile.getName() + " between "
									+ prevFrame + " " + frame);
						
						// Second column contains the gestures
						String[] gesturesTokens = tokens[1].replaceAll("[\\]\\[]", "").trim().split(",");
						for (int i = 0; i < gesturesTokens.length; i++)
							if (!"".equals(gesturesTokens[i].trim()))
								gestures.add(gesturesTokens[i].trim());
					}

					// Now that we parsed the targets, it's time to add to the intervals.
					// Note: We will need to interpolate between any missing frames for now
					// 		since we're working with spotty data.
					// For each rolling interval, we'll either extend it if the target is
					// still valid, or end it otherwise.
					for (int i = currentIntervals.size()-1; i >= 0; i--) {
						ParamInterval<String> interval = currentIntervals.get(i);
						if (gestures.contains(interval.obj())) {
							// Extend the interval if the gesture is still valid.
							interval.interval().setEnd(frame+1);
							gestures.remove(interval.obj());
						} else {
							// Otherwise, end it.
//							interval.interval().setEnd(frame);
							actor.addGestureInterval(interval);
							currentIntervals.remove(i);
						}
					}
					
					// For each new target, start a new rolling interval
					for (String gesture : gestures) {
						assert (gesture != null && !"".equals(gesture)) : "Invalid gesture!";
						currentIntervals.add(new ParamInterval<String>(frame, frame+1, gesture));
					}
					gestures.clear();
					
					prevFrame = frame;
				}
				
				// By the time we finish reading the file, if there are any rolling intervals
				// left, then end them.
				for (int i = currentIntervals.size()-1; i >= 0; i--) {
					ParamInterval<String> interval = currentIntervals.get(i);
					interval.interval().setEnd(frame+1);
					actor.addGestureInterval(interval);
					currentIntervals.remove(i);
				}
			}
			
		} catch (Exception e) {
			System.out.println("Exception while reading file: " + e);
		}

		// Dump to file
		if (out) {
			StringBuilder sb = new StringBuilder();
			for (Person p : this.getActors())
				sb.append(p.toStringGestures() + "\n");
			System.out.println(sb.toString());
			Utils.toFile(sb, Utils.OUTDIR + _name + "-gestures.txt", false);
		}
	}
	
	/**
	 * Load pairs-probabilities data from file.
	 * @param dataFile Pairs data file.
	 */
	public void loadPairsData(File dataFile) {
		try {
			// Create BufferedReader to read csv file
			BufferedReader br = new BufferedReader(new FileReader(dataFile));
			String strLine = "";
			int prevFrame = -1, frame = 0;
			final String UNKNOWN_VALUE = "-1";
			
			// First line gives us pair combinations
			strLine = br.readLine();
			String[] tokens = strLine.trim().split(",");
			List<Set<Person>> pairKeys = new ArrayList<Set<Person>>();
			for (int i = 1; i < tokens.length; i++) {
				// First column is "frame", subsequent columns give us
				// the pair combinations.
				String[] pairTokens = tokens[i].trim().split(" ");
				Set<Person> pair = new HashSet<Person>();
				
				// Note: Pairs can now have any arbitrary number of members
				for (String mem : pairTokens) {
					ActorName p = Utils.getActorNameFromTrackID(_name, mem);
					pair.add(this.getOrAddPerson(p));
				}
				
				pairKeys.add(pair);
				_pairs.put(pair, new ArrayList<Double>());
			}

			// Read comma separated file line by line
			while ((strLine = br.readLine()) != null) {

				// Break comma separated line using ","
				tokens = strLine.split(",");
				if (tokens.length > 1) {
					
					// First column is frame number
					frame = Integer.parseInt(tokens[0].trim());
					if (prevFrame != -1 && frame != prevFrame+1)
						System.err.println("Missing pairs frame(s) in "
							+ dataFile.getName() + " between "
							+ prevFrame + " " + frame);
					
					// Subsequent columns are pairs probabilities
					for (int i = 1; i < tokens.length; i++) {
						Set<Person> pair = pairKeys.get(i-1);
						if (tokens[i].trim().equals(UNKNOWN_VALUE))
							_pairs.get(pair).add(-1.0);
						else
							_pairs.get(pair).add(Double.parseDouble(tokens[i].trim()));
					}
				}
				
				prevFrame = frame;
			}
			
			// Sanity check here
			int numFrames = -1;
			for (List<Double> l : _pairs.values()) {
				assert numFrames == -1 || l.size() == numFrames :
					"Uneven pairs distributions";
				numFrames = l.size();
			}

		} catch (Exception e) {
			System.out.println("Exception while reading csv file: " + e);
		}
	}
	
	
	/**
	 * Print the grouping information of the scene in a nice format.
	 * @return The grouping description of the scene.
	 */
	public StringBuilder printGroups() {
		StringBuilder sb = new StringBuilder();
		Collections.sort(_groups);
		sb.append("\n== " + _groups.size() + " Groups ==\n\n");
		for (Group gr : _groups)
			sb.append(gr.toString() + "\n");
		System.out.println(sb.toString());
		return sb;
	}
	
}
